package com.kara.woodAgent.agent.wood.node;

import com.alibaba.fastjson2.JSON;
import com.kara.woodAgent.agent.graph.Next;
import com.kara.woodAgent.agent.graph.Node;
import com.kara.woodAgent.agent.model.ModelProvider;
import com.kara.woodAgent.agent.wood.context.WoodContext;
import com.kara.woodAgent.agent.wood.model.FacadeResult;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ResponseFormat;
import dev.langchain4j.model.chat.request.ResponseFormatType;
import dev.langchain4j.model.chat.response.ChatResponse;
import org.springframework.beans.factory.annotation.Autowired;

import java.util.List;

public class FacadeNode implements Node<WoodContext> {


    private final String systemPrompt = """
                 #### 身份
                 你是wood智能体系统 是由 wood团队开发的通用智能体系统.
                 #### 职责
                 你的任务是判断用户输入的内容是任务还是对话，如果是是任务就需要调用其他智能体来解决,如果是对话就直接回答。
                 #### 返回结果
                 返回结果为json格式,包含以下两个字段:
                 - type: 类型,可以是 direct 或者 dispatch, 如果用户输入可以直接回答,则返回 direct,如果用户输入需要调用其他智能体,则返回 dispatch (必填)
                 - message: 当type为direct时不为空,内容为对于用户输入的回答.
            """;
    @Autowired
    private ModelProvider modelProvider;



    @Override
    public Next execute(WoodContext context) {

        ChatLanguageModel chatModel = modelProvider.getChatModel();
        String userQuestion = context.getUserQuestion();
        List<ChatMessage> messages = List.of(
                SystemMessage.from(systemPrompt),
                UserMessage.from(userQuestion)
        );

        ChatRequest request = ChatRequest.builder().messages(messages).responseFormat(ResponseFormat.builder().type(ResponseFormatType.JSON).build())
                .build();
        ChatResponse chatResponse = chatModel.chat(request);

        AiMessage aiMessage = chatResponse.aiMessage();
        String text = aiMessage.text();

        FacadeResult facadeResult = JSON.parseObject(text, FacadeResult.class);

        if ("direct".equals(facadeResult.type())) {
            return Next.End().step(text);
        } else {
            return Next.NextNode("planning").step(text);
        }
    }
}
